//ffmpeg.cpp
#include "ffmpeg.h"

AVFormatContext * m_in_format_ctx = NULL;
int video_stream_idx = -1;

void initialize_all()
{
    av_register_all();
    avformat_network_init();
}

int init_input(char * filename,AVFormatContext ** in_format_ctx)
{
    int ret = 0;
    ret = avformat_open_input(in_format_ctx, filename,NULL, NULL);
    if (ret != 0)
    {
        printf("Call avformat_open_input function failed!\n");
        return 0;
    }
    if (av_find_stream_info(*in_format_ctx) < 0)
    {
        printf("Call av_find_stream_info function failed!\n");
        return 0;
    }

    ret = open_codec_context(&video_stream_idx, in_format_ctx, AVMEDIA_TYPE_VIDEO);
    if (ret != 1)
    {
        printf("Call open_codec_context function failed!\n");
        return 0;
    }

    //The output of input information
    av_dump_format(*in_format_ctx, -1, filename, 0);

    return 1;
}

int open_codec_context(int *stream_idx, AVFormatContext ** in_format_ctx, enum AVMediaType type)
{
    int ret = 0;
    AVStream * st = NULL;
    AVCodecContext * dec_ctx = NULL;
    AVCodec * dec = NULL;

    ret = av_find_best_stream((*in_format_ctx), type, -1, -1, NULL, 0);
    if (ret < 0)
    {
        printf("Call av_find_best_stream function failed!\n");
        return 0;
    }
    else
    {
        *stream_idx = ret;
        st = (*in_format_ctx)->streams[*stream_idx];

        /* find decoder for the stream */
        dec_ctx = st->codec;
        dec = avcodec_find_decoder(dec_ctx->codec_id);
        if (!dec)
        {
            return AVERROR(EINVAL);
        }

        if ((ret = avcodec_open2(dec_ctx, dec, NULL)) < 0)
        {
            printf("Call avcodec_open2 function failed!\n");
            return 0;
        }
    }
    return 1;
}

void uinit_input(AVFormatContext * in_format_ctx)
{
    avcodec_close(in_format_ctx->streams[video_stream_idx]->codec);
    avformat_close_input(&in_format_ctx);
    av_free(in_format_ctx);
}

int read_frame(AVFormatContext * in_format,char * filepath)
{
    int ret = 0;
    AVPacket pkt_in;
    AVPacket pkt_out;
    AVFrame * picture = av_frame_alloc();
    int got_picture_ptr = 0;
    AVFormatContext* pFormatCtx = NULL;
    AVStream* video_st = NULL;
    AVCodecContext* pCodecCtx = NULL;
    AVCodec* pCodec = NULL;
    FILE * video_dst_file = NULL;
    char file_name[50] = {0};

    /* find the mpeg1 video encoder */
    pCodec = avcodec_find_encoder(AV_CODEC_ID_MJPEG);
    if (!pCodec)
    {
        printf("Call avcodec_find_encoder function failed!\n");
        return 0;
    }

    pCodecCtx = avcodec_alloc_context3(pCodec);
    if (!pCodecCtx)
    {
        printf("Call avcodec_alloc_context3 function failed!\n");
        return 0;
    }


    pCodecCtx->codec_type = AVMEDIA_TYPE_VIDEO;
    pCodecCtx->pix_fmt = AV_PIX_FMT_YUVJ420P;

    //If you want to change the image size can do the zoom function
    pCodecCtx->width = in_format->streams[video_stream_idx]->codec->width;
    pCodecCtx->height = in_format->streams[video_stream_idx]->codec->height;

    pCodecCtx->time_base.num = 1;
    pCodecCtx->time_base.den = 15;
    //Output some information

    if (avcodec_open2(pCodecCtx, pCodec,NULL) < 0)
    {
        printf("Could not open codec.\n");
        return 0;
    }

    av_init_packet(&pkt_in);
    av_init_packet(&pkt_out);

    //return 0 if OK, < 0 on error or end of file
    while (av_read_frame(in_format, &pkt_in) == 0)
    {
        //video && key frame
        if (pkt_in.stream_index == video_stream_idx && pkt_in.flags == 1)
        {
            ret = avcodec_decode_video2(in_format->streams[video_stream_idx]->codec, picture, &got_picture_ptr, &pkt_in);
            //success decode
            if(got_picture_ptr)
            {
                av_new_packet(&pkt_out,pCodecCtx->width * pCodecCtx->height * 3);
                //encode
                ret = avcodec_encode_video2(pCodecCtx, &pkt_out,picture, &got_picture_ptr);
                if(ret < 0)
                {
                    printf("Encode Error.\n");
                    avcodec_close(pCodecCtx);
                    av_free(pCodecCtx);
                    av_free_packet(&pkt_out);
                    return 0;
                }
                //success encode
                if (got_picture_ptr)
                {
                    generate_file_name(file_name,filepath,pkt_in.pts);

                    //Method 1
                    //save_picture_init(&pFormatCtx,video_st,pCodecCtx,file_name);
                    //save_picture_uinit_2(pFormatCtx,pkt_out);
                    //Method 2
                    save_picture_init_2(&video_dst_file,file_name);
                    if(video_dst_file)
                    {
                        save_picture_uinit(video_dst_file,pkt_out);
                    }
                }
            }
        }
        av_free_packet(&pkt_in);
    }
    av_frame_free(&picture);
    avcodec_close(pCodecCtx);
    av_free(pCodecCtx);

    return 1;
}

int  save_picture_init(AVFormatContext** out_format_ctx,AVStream * video_st,AVCodecContext* pCodecCtx,char * filename)
{
    AVOutputFormat* fmt = NULL;

    avformat_alloc_output_context2(out_format_ctx, NULL, NULL, filename);
    fmt = (*out_format_ctx)->oformat;

    video_st = avformat_new_stream((*out_format_ctx), 0);
    if (video_st == NULL)
    {
        return 0;
    }
    /*pCodecCtx = video_st->codec;
    pCodecCtx->codec_id = fmt->video_codec; */
    av_dump_format((*out_format_ctx), 0, filename, 1);

    if (avformat_write_header((*out_format_ctx), NULL))
    {
        printf("Call avformat_write_header function failed.\n");
        return 0 ;
    }
    return 1;
}

int  save_picture_init_2(FILE ** pFile,char * filename)
{
    *pFile = fopen(filename, "wb");
    return 1;
}

void  save_picture_uinit(FILE * pFile,AVPacket pkt)
{
    fwrite(pkt.data, sizeof(uint8_t),pkt.size, pFile);
    fclose(pFile);
    av_free_packet(&pkt);
}

void  save_picture_uinit_2(AVFormatContext* out_format_ct,AVPacket pkt)
{
    int ret = 0;

    //pkt.stream_index = video_st->index;
    ret = av_interleaved_write_frame(out_format_ct, &pkt);
    if (ret != 0)
    {
        printf("error av_interleaved_write_frame _ video\n");
        return ;
    }
    //Write Trailer
    ret = av_write_trailer(out_format_ct);
    if (ret != 0)
    {
        printf("error av_interleaved_write_frame _ video\n");
    }
    av_free(out_format_ct);
}

void generate_file_name(char * file_name,char * file_path,long long pts)
{
    if (file_path[strlen(file_path) - 1] == '/')
    {
        sprintf(file_name,"%s%lld.jpg",file_path,pts);
    }
    else
    {
        sprintf(file_name,"%s/%lld.jpg",file_path,pts);
    }
}